package uk.ac.aber.paws.client.net;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.StringTokenizer;

import org.xbill.DNS.AAAARecord;
import org.xbill.DNS.ARecord;
import org.xbill.DNS.ExtendedResolver;
import org.xbill.DNS.Lookup;
import org.xbill.DNS.Name;
import org.xbill.DNS.PTRRecord;
import org.xbill.DNS.Record;
import org.xbill.DNS.Resolver;
import org.xbill.DNS.ReverseMap;
import org.xbill.DNS.TextParseException;
import org.xbill.DNS.Type;

/**
 * @author Rob Johnson, Aberystwyth University
 * @email rob.johnson@aber.ac.uk
 * @date 16th August 2011
 * 
 * DNSLookup.java
 * 
 * Resolves a host name based off the computer's IP address - the equivalent
 * of nslookup for example. It is used on PAWS Client initialisation to determine
 * the host name of the client. This valuable parameter is used as a basis of 
 * logging identity to the database. This code has dependancy on the org.xbill.DNS
 * DNSJava library (dnsjava-2.1.1.jar) www.dnsjava.org , http://sourceforge.net/projects/dnsjava/
 * 
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

public class DNSLookup implements InvocationHandler {

	private static final String nsProperty = "sun.net.spi.nameservice.nameservers";
	private static final String domainProperty = "sun.net.spi.nameservice.domain";
	private static final String v6Property = "java.net.preferIPv6Addresses";

	private boolean preferV6 = false;

	public DNSLookup() {
		
		String nameServers = System.getProperty(nsProperty);
		String domain = System.getProperty(domainProperty);
		String v6 = System.getProperty(v6Property);
	
		if (nameServers != null) {
			
			StringTokenizer st = new StringTokenizer(nameServers, ",");
			String [] servers = new String[st.countTokens()];
			int n = 0;
			while (st.hasMoreTokens())
				servers[n++] = st.nextToken();
			try {
				Resolver res = new ExtendedResolver(servers);
				Lookup.setDefaultResolver(res);
			}
			catch (UnknownHostException e) {
				System.err.println("DNSJavaNameService: invalid " +
						   nsProperty);
			}
		}
	
		if (domain != null) {
			try {
				Lookup.setDefaultSearchPath(new String[] {domain});
			}
			catch (TextParseException e) {
				System.err.println("DNSJavaNameService: invalid " +
						   domainProperty);
			}
		}
	
		if (v6 != null && v6.equalsIgnoreCase("true"))
			preferV6 = true;
	}
	
	public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
		try {
			if (method.getName().equals("getHostByAddr")) {
				return this.getHostByAddr((byte[]) args[0]);
			} 
			else if (method.getName().equals("lookupAllHostAddr")) {
				
				InetAddress[] addresses = this.lookupAllHostAddr((String) args[0]);
				
				if (method.getReturnType().equals(InetAddress[].class)) {
					
					// method for Java >= 1.6
					return addresses;
				} 
				else if (method.getReturnType().equals(byte[][].class)) {
					// method for Java <= 1.5
					byte[][] byteAddresses = new byte[addresses.length][];
					
					for (int i=0; i < addresses.length; i++) {
						byteAddresses[i] = addresses[i].getAddress();
					}
					return byteAddresses;
				}
			}		
		} 
		catch (Throwable e) {
			System.err.println("DNSJavaNameService: Unexpected error.");
			e.printStackTrace();
			throw e;
		}
		throw new IllegalArgumentException("Unknown function name or arguments.");
	}
	
	/**
	 * Performs a forward DNS lookup for the host name.
	 * @param host The host name to resolve.
	 * @return All the ip addresses found for the host name.
	 */
	public InetAddress[] lookupAllHostAddr(String host) throws UnknownHostException {
		
		Name name = null;
	
		try {
			name = new Name(host);
		}
		catch (TextParseException e) {
			throw new UnknownHostException(host);
		}
	
		Record [] records = null;
		if (preferV6)
			records = new Lookup(name, Type.AAAA).run();
		if (records == null)
			records = new Lookup(name, Type.A).run();
		if (records == null && !preferV6)
			records = new Lookup(name, Type.AAAA).run();
		if (records == null)
			throw new UnknownHostException(host);
	
		InetAddress[] array = new InetAddress[records.length];
		
		for (int i = 0; i < records.length; i++) {
			Record record = records[i];
			if (records[i] instanceof ARecord) {
				ARecord a = (ARecord) records[i];
				array[i] = a.getAddress();
			} else {
				AAAARecord aaaa = (AAAARecord) records[i];
				array[i] = aaaa.getAddress();
			}
		}
		return array;
	}
	
	/**
	 * Performs a reverse DNS lookup.
	 * @param addr The ip address to lookup.
	 * @return The host name found for the ip address.
	 */
	public String getHostByAddr(byte [] addr) throws UnknownHostException {
		
		Name name = ReverseMap.fromAddress(InetAddress.getByAddress(addr));
		Record [] records = new Lookup(name, Type.PTR).run();
		
		if (records == null)
			throw new UnknownHostException();
		
		String hostName = ((PTRRecord) records[0]).getTarget().toString();
		
		// remove the trailing "."
		if (hostName.toCharArray()[hostName.length()-1] == '.') {
			return hostName.substring(0, hostName.length()-1);
		}
		
		return hostName;
	}
}